# 引入相关的库
import random
import math
import pandas as pd
import numpy as np

# 模型类
class LFM(object):

    """
    :param rating_data: rating_data是[(user,[(item,rate)]]类型
    :param F: 隐因子个数
    :param alpha: 学习率
    :param lmbd: 正则化
    :param max_iter:最大迭代次数
    """
    def __init__(self, F,file_path,alpha=0.01, lmbd=0.01, max_iter=100):
        self.F = F
        self.P = dict()  # R=PQ^T，代码中的Q相当于博客中Q的转置
        self.Q = dict()
        self.alpha = alpha
        self.lmbd = lmbd
        self.max_iter = max_iter
        self.file_path = file_path
        self.rating_data = self.get_rating_data()

        '''随机初始化矩阵P和Q'''
        for user, rates in self.rating_data.items():
            self.P[user] = [random.random() / math.sqrt(self.F) for x in range(self.F)]
            for item, _ in rates.items():
                if item not in self.Q:
                    self.Q[item] = [random.random() / math.sqrt(self.F) for x in range(self.F)]

    def get_rating_data(self):
        rating_data = dict()
        # 训练数据集
        with open(self.file_path, "r") as f:
            for line in f.readlines()[1:]:
                user_id, item_id, score = line.split(",")
                if float(score) >= 1.0:
                    score = 3.5
                else:
                    score = 1.5
                #print(float(score))
                rating_data.setdefault(user_id,dict())
                rating_data[user_id][item_id] = score

        return rating_data


    """
    随机梯度下降法训练参数P和Q
    :return: 
    """
    def train(self):
        for step in range(self.max_iter):
            for user, rates in self.rating_data.items():
                for item, rui in rates.items():
                    hat_rui = self.predict(user, item)
                    err_ui = rui - hat_rui
                    for f in range(self.F):
                        self.P[user][f] += self.alpha * (err_ui * self.Q[item][f] - self.lmbd * self.P[user][f])
                        self.Q[item][f] += self.alpha * (err_ui * self.P[user][f] - self.lmbd * self.Q[item][f])
            self.alpha *= 0.9
            
            print("{}轮".format(step))

    """
    :param user:
    :param item:
    :return:
    预测用户user对物品item的评分
    """
    def predict(self, user, item):
        return sum(self.P[user][f] * self.Q[item][f] for f in range(self.F))

    def predict_test(self):
        result = []
        for user,item in self.test:
            try:
                rate = self.predict(user,item)
            except:
                rate = 2.75

            if rate>2.5:
                result.append(1)
            else:
                result.append(0)

        save_csv(result,"new_one")


def save_csv(cnt, name):
    ids = [i for i in range(1, len(cnt) + 1)]
    df_data = []
    for i in range(len(ids)):
        temp = []
        temp.append(ids[i])
        temp.append(cnt[i])
        df_data.append(temp)

    df = pd.DataFrame(np.array(df_data), columns=['id', 'rating'])
    df.to_csv("temp/{}.csv".format(name), index=False)
    print("{} 写入完成".format(name))
    
def get_test_data(pt):
    test = list()
    with open(pt,"r") as f:
        for line in f.readlines()[1:]:
            user_id, item_id = line[:-1].split(",")
            test.append((user_id, item_id))
    test = np.array(test)
    return test
    
if __name__ =="__main__":
    p0 = "dataSets/train_0.csv"
    p1 = "dataSets/train_1.csv"
    p2 = "dataSets/train_2.csv"
    pt = "dataSets/test.csv"

    test = get_test_data(pt)

    print(test)
    lfm0 = LFM(F=60,file_path=p0)
    lfm1 = LFM(F=60,file_path=p1)
    lfm2 = LFM(F=60,file_path=p2)
    
    # 模型训练
    lfm0.train()
    lfm1.train()
    lfm2.train()
    
    # 结果保存
    r0,r00 = [],[]
    r1,r11 = [],[]
    r2,r22 = [],[]
    
    for user,item in test:
        try:
            rate = lfm0.predict(user,item)
            r00.append(rate)
        except:
            rate = 2.75
            r00.append(2.5)
            
        if rate>=2.5:
            r0.append(1)
        else:
            r0.append(0)
            
    save_csv(r0,"r0")
    save_csv(r00,"r00")
    

    for user,item in test:
        try:
            rate = lfm1.predict(user,item)
            r11.append(rate)
        except:
            rate = 2.75
            r11.append(2.5)

        if rate>=2.5:
            r1.append(1)
        else:
            r1.append(0)

    save_csv(r1,"r1")
    save_csv(r11,"r11")


    for user,item in test:
        try:
            rate = lfm2.predict(user,item)
            r22.append(rate)
        except:
            rate = 2.75
            r22.append(2.5)

        if rate>=2.5:
            r2.append(1)
        else:
            r2.append(0)

    save_csv(r2,"r2")
    save_csv(r22,"r22")


    # 蓉儿模型
    r_finall = []
    for i in range(len(r00)):
        if ((r00[i]+r11[i]+r22[i])/3 >= 2.5):
            r_finall.append(1)
        else:
            r_finall.append(0)

    save_csv(r_finall,"finally")
    